# -----------------------------------------------------------------------
from idaapi import *
from idc import *
from idautils import *
# This script detects a loop in the current function.
# (c) Hex-Rays
#

class LoopDetector():
    def __init__(self, location):
        self.location = location
        self.loopSegments = []
        self.nestedLoopSegments = []

    def DetectLoops(self):
               
        f = idaapi.FlowChart(idaapi.get_func(here()))
        self.nodes = []
        self.backedges = []
        self.time = 0
        for block in f:
            self.nodes.append(Node(Block=block, Color="WHITE", Start=0, Finish=0))        

        for node in self.nodes:
            if node.Color == "WHITE":
                self.dfsVisit(node)

        for loop in self.backedges:
            #SetColor(loop.To.startEA, CIC_ITEM, 0xa07420)
            #SetColor(loop.From.endEA, CIC_ITEM, 0xa07420)
            self.loopSegments.append([loop.To.startEA, loop.From.endEA])
            
        return int(len(self.backedges))

    def dfsVisit(self, node):
        node.Color = "GRAY"
        self.time = self.time+1
        node.Start = self.time
        
        for adj in self.nodes:            
            if [True for succs in node.Block.succs() if adj.Block.id == succs.id]:
                if adj.Color == "WHITE":
                    self.dfsVisit(adj)
                elif adj.Color == "GRAY":                    
                    self.backedges.append(BackEdge(From=node.Block, To=adj.Block))
        node.Color = "BLACK"
        self.time = self.time+1
        node.Finish = self.time
        return

    def DetectNestedLoops(self):
        
        if len(self.loopSegments) == 0:
            self.DetectLoop()
            if len(self.loopSegments) == 0:
                return 0

        self.loopSegments.sort()
      
        for i in range(len(self.loopSegments)-1):
            j = i + 1
            loopOuter = self.loopSegments[i]
            loopInner = self.loopSegments[j]
            
            if loopOuter[1] > loopInner[0] and loopOuter[1] > loopInner[1] :
                #SetColor(loopOuter[0], CIC_ITEM, 0xd44444)
                #SetColor(loopOuter[1], CIC_ITEM, 0xd44444)
                self.nestedLoopSegments.append([loopOuter[0], loopOuter[1], loopInner[0], loopInner[1]]) 

        return int(len(self.nestedLoopSegments))

    def DetectInstruction(self, insn):
                
        hasComparison = False
        nOccurences = 0
        for nest in self.nestedLoopSegments:
            for x in [x for x in Heads(nest[2], nest[3])]:
               
                instruction = idc.GetDisasm(x).split(' ')[0]              
                if instruction == insn:
                    nOccurences = nOccurences + 1
        return nOccurences

    def DetectComparisons(self):
        return self.DetectInstruction("cmp")

def detect_features(seg):
    detLoops = LoopDetector(seg)
    loops = detLoops.DetectLoops()
    nest_loops = 0
    comparisons = 0
    if loops > 0:
        nest_loops = detLoops.DetectNestedLoops()
        comparisons = detLoops.DetectComparisons()
    return [loops, nest_loops, comparisons]

locals()['ret'] = detect_features(loc)